import enum
import math
import numpy as np
import scipy.special
import python.version as version

verbose = True


def set_verbose(v):
    global verbose
    verbose = v


def print_v(*args, end=None):
    return 0
    # global verbose
    # if verbose:
    #     print(*args, end=end)


class IneqType(enum.Enum):
    LE = (True,)
    GE = (False,)
    DB = ()  # 双边不等式


def check_inequalities(key, ineqs):
    for ineq in ineqs:
        # for ineq, veci, sign, _, is_correct in ineqs:
        if not ineq.is_correct:
            continue
        dot = sum([a * b for a, b in zip(ineq.coefficients, key)])
        if ineq.sign == IneqType.DB:
            if (dot < ineq.b[0]) or (dot > ineq.b[1]):
                print(ineq.b, dot)
                return False
        elif ineq.sign == IneqType.LE:
            if dot > ineq.b[0]:
                return False
        else:  # ineq.sign == IneqType.GE:
            if dot < ineq.b[0]:
                return False
    return True


def bino(eta):
    binp = lambda x: scipy.special.binom(eta, x) / (2**eta)
    binom1 = [binp(k) for k in range(eta + 1)]
    e = np.convolve(binom1, binom1)
    dist = {
        i: 0.0 for i in range(-version.KYBER_k * eta, version.KYBER_k * eta + 1)
    }  # 初始化所有键
    for i in range(len(e)):
        dist[i - eta] = e[i].item()
    return dist


def evol_bino(eta):
    binp = lambda x: scipy.special.binom(eta, x) / (2**eta)
    binom1 = [binp(k) for k in range(eta + 1)]
    e = np.convolve(binom1, binom1)
    evol = e
    for _ in range(version.KYBER_k - 1):
        evol = np.convolve(evol, e)
    return {i - version.KYBER_k * eta: evol[i] for i in range(len(evol))}


def most_likely(props_dicts):
    # 80: ({1: 0.25, 0: 0.375, 2: 0.0625, -2: 0.0625, -1: 0.25}, 2.0306390622295662) 返回概率最大的键值 0
    res = [0] * len(props_dicts)
    for d in props_dicts:
        res[d] = max(props_dicts[d][0], key=lambda v: props_dicts[d][0][v])
    return res


def most_likely_coeff(dist):
    return max(dist, key=dist.get)


def most_likely_list(probs_list):
    # 与most_likely不同传入参数为[{1: 0.25, 0: 0.375, 2: 0.0625, -2: 0.0625, -1: 0.25},...]
    return [most_likely_coeff(d) for d in probs_list]


def expected(dist):
    # 分布的期望
    return sum((i * p for i, p in dist.items()))


def var(dist):
    # 计算分布的方差
    exp = expected(dist)
    return sum((p * (exp - i) ** 2 for i, p in dist.items()))


def mat_mul(mat, vec, q=None):
    # 矩阵与向量乘法
    if q is None:
        return (sum((m_ij * v_j for m_ij, v_j in zip(m_i, vec))) for m_i in mat)
    else:
        return (
            sum(((m_ij * v_j) % q for m_ij, v_j in zip(m_i, vec))) % q for m_i in mat
        )


def euclidean_dist(b0, b1):
    # 欧几里得距离
    assert len(b0) == len(b1)
    return math.sqrt(sum(((b0i - b1i) ** 2) for b0i, b1i in zip(b0, b1)))


def add_vec(v0, v1, q=None):
    if q is None:
        return map(lambda x: x[0] + x[1], zip(v0, v1))
    else:
        return map(lambda x: (x[0] + x[1]) % q, zip(v0, v1))


def sub_vec(v0, v1, q=None):
    # 向量减法
    if q is None:
        return map(lambda x: x[0] - x[1], zip(v0, v1))
    else:
        return map(lambda x: (x[0] - x[1]) % q, zip(v0, v1))


def dot(a, b, q=None):
    if q:
        return sum((ai * bi % q for ai, bi in zip(a, b)))
    else:
        return sum((ai * bi for ai, bi in zip(a, b)))


def norm(a):
    # 二范数
    return math.sqrt(dot(a, a))


def reduce_sym(a, q=3329):
    # 求中心模
    a %= q
    if a > q // 2:
        a -= q
    if a <= -q // 2:
        a += q
    return a


def reduce_sym_list(a, q=3329):
    return list(map(lambda x: reduce_sym(x, q), a))


def mod_inverse(a, q=3329):
    # 使用扩展的欧几里得算法计算乘法逆元
    t, new_t = 0, 1
    r, new_r = q, a

    while new_r != 0:
        quotient = r // new_r
        t, new_t = new_t, t - quotient * new_t
        r, new_r = new_r, r - quotient * new_r

    if r > 1:
        raise ValueError("a is not invertible")

    if t < 0:
        t += q

    return t
